# ====================================================================================
# 文件: src/utils.py
# 描述: [V_Final] 新增 info_nce 实现，用于 SRA-CL 模块
# ====================================================================================

import logging
import os
import random
import numpy as np
import scipy.sparse as sp
import torch
import torch.nn.functional as F


def get_sparse_adj(adj_matrix, device, gmae_norm=False):
    """ (保持原有的 V30 修复版不变) """
    sp_adj = adj_matrix.tocoo()
    shape = sp_adj.shape
    row_A, col_A = sp_adj.row, sp_adj.col
    values_A = sp_adj.data.astype(np.float32)
    row_I = np.arange(shape[0])
    col_I = np.arange(shape[0])
    values_I = np.ones(shape[0], dtype=np.float32)
    row_AI = np.concatenate([row_A, row_I])
    col_AI = np.concatenate([col_A, col_I])
    values_AI = np.concatenate([values_A, values_I])

    adj_with_selfloop_sp = sp.csr_matrix((values_AI, (row_AI, col_AI)), shape=shape)
    adj_with_selfloop_sp.sum_duplicates()
    adj_coo = adj_with_selfloop_sp.tocoo()
    row, col = adj_coo.row, adj_coo.col
    values = adj_coo.data

    deg_row = np.array(adj_with_selfloop_sp.sum(axis=1)).flatten()
    if not gmae_norm:
        deg_col = np.array(adj_with_selfloop_sp.sum(axis=0)).flatten()

    if gmae_norm:
        with np.errstate(divide='ignore'):
            d_inv = np.power(deg_row, -1)
        d_inv[np.isinf(d_inv)] = 0.
        norm_coeffs = d_inv[row]
    else:
        with np.errstate(divide='ignore'):
            d_inv_sqrt_row = np.power(deg_row, -0.5)
            d_inv_sqrt_col = np.power(deg_col, -0.5)
        d_inv_sqrt_row[np.isinf(d_inv_sqrt_row)] = 0.
        d_inv_sqrt_col[np.isinf(d_inv_sqrt_col)] = 0.
        norm_coeffs = d_inv_sqrt_row[row] * d_inv_sqrt_col[col]

    norm_weights_cpu = (norm_coeffs * values).astype(np.float32)
    edge_index_cpu = torch.LongTensor(np.vstack((row, col)))
    edge_weight_cpu = torch.FloatTensor(norm_weights_cpu)

    return edge_index_cpu.to(device), edge_weight_cpu.to(device)


# [新增] InfoNCE 损失函数 (用于 SRA-CL)
def info_nce_loss(features_1, features_2, temperature=0.1):
    """
    计算两个视图间的 InfoNCE 损失 (Contrastive Loss)
    features_1: [batch_size, dim]
    features_2: [batch_size, dim]
    """
    # L2 归一化
    features_1 = F.normalize(features_1, dim=1)
    features_2 = F.normalize(features_2, dim=1)

    batch_size = features_1.shape[0]

    # 计算相似度矩阵: [batch_size, batch_size]
    logits = torch.matmul(features_1, features_2.T) / temperature

    # 正样本是由于同一物品在不同视图下的表示 (对角线)
    labels = torch.arange(batch_size, device=features_1.device, dtype=torch.long)

    # 交叉熵损失
    loss = F.cross_entropy(logits, labels)
    return loss


# 基础工具
def set_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def get_logger(log_file: str):
    logger = logging.getLogger(f"SDKR::{log_file}")
    if logger.handlers: return logger
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter("%(asctime)s | %(levelname)s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(formatter)
    logger.addHandler(stream_handler)
    os.makedirs(os.path.dirname(log_file), exist_ok=True)
    file_handler = logging.FileHandler(log_file)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    logger.propagate = False
    return logger


# 评估指标
def recall_at_k(true_list, pred_list, k):
    relevant_items = set(true_list)
    predicted_items = set(pred_list[:k])
    if not relevant_items: return 0.0
    return len(relevant_items.intersection(predicted_items)) / len(relevant_items)


def ndcg_at_k(true_list, pred_list, k):
    relevant_items = set(true_list)
    if not relevant_items: return 0.0
    dcg = 0.0
    for i, item_id in enumerate(pred_list[:k]):
        if item_id in relevant_items: dcg += 1.0 / np.log2(i + 2)
    ideal_k = min(k, len(relevant_items))
    idcg = sum([1.0 / np.log2(i + 2) for i in range(ideal_k)])
    return dcg / idcg if idcg > 0.0 else 0.0


def precision_at_k(true_list, pred_list, k):
    relevant_items = set(true_list)
    predicted_items = set(pred_list[:k])
    if not predicted_items: return 0.0
    return len(relevant_items.intersection(predicted_items)) / len(predicted_items)


def hit_ratio_at_k(true_list, pred_list, k):
    relevant_items = set(true_list)
    predicted_items = set(pred_list[:k])
    if not relevant_items: return 0.0
    return 1.0 if len(relevant_items.intersection(predicted_items)) > 0 else 0.0
